# -*-coding=utf-8-*-


import os
import re
import numpy as np
import time


class PES():
    def __init__(self):
        self.lines = []
        # 当前文件夹绝对路径
        self.path = os.path.abspath('.')
        self.log_path = self.path + '\\log.txt'
        self.opt = '# p force b3lyp/6-31g'
        self.head_pattern = re.compile(r' \*{5} Axes restored to original set \*{5}')
        self.data_pattern = re.compile(r'\s+?\d+?\s+?(\d+?)\s+?(-?\d+\.?\d*?)\s+(-?\d+\.?\d*?)\s+(-?\d+\.?\d*?)\s+')
        self.energy_pattern = re.compile(r' SCF Done:  E\(RB3LYP\) =  (-?\d+\.?\d*?)\s')
        self.normal_end_pattern = re.compile(r' Normal termination of Gaussian 09 at')
        self.element = []
        self.position = None
        self.n = 0
        self.file_number = 1

    def set_path(self, path):
        """
        设置路径
        """
        self.path = path
        self.log_path = path + '\\' + 'log.txt'

    def load_file(self, path):
        """
        行模式读取文件文本
        """
        with open(path, 'r') as f:
            self.lines = f.readlines()
        if self.check_normal_end(self.lines[-1]):
            self.write_log('Normal termination :' + path)
        else:
            self.write_log('Error termination :' + path)
            exit()

    def search_force(self):
        """
        搜索梯度
        """
        force = []
        head_line = 0
        for i in range(len(self.lines)):
            result = re.match(self.head_pattern, self.lines[i])
            if result:
                head_line = i
                break
        data_start_line = head_line + 5
        for i in range(data_start_line, len(self.lines)):
            result = re.match(self.data_pattern, self.lines[i])
            if result:
                for item in result.group(2, 3, 4):
                    # 力单位: eV/A
                    force.append(float(item) * 51.4224725)
        return force

    def search_energy(self):
        """
        搜索能量
        """
        energy = 0
        for i in range(len(self.lines)):
            result = re.match(self.energy_pattern, self.lines[i])
            if result:
                # 能量单位: kcal/mol
                energy = float(result.group(1)) * 627.5095
                break
        return energy

    def init_load(self, path):
        """
        初始读取文件
        """
        pattern = re.compile(r' (\w+?)\s+?(-?\d+\.?\d*)\s+?(-?\d+\.?\d*)\s+?(-?\d+\.?\d*)')
        lines = []
        with open(path, 'r') as f:
            lines = f.readlines()
        start_read = False
        position = []
        for i in range(len(lines)):
            result = re.match(pattern, lines[i])
            if result:
                start_read = True
                # 元素名
                self.element.append(result.group(1))
                # 坐标
                for item in result.group(2, 3, 4):
                    position.append(float(item))
                if i == len(lines)-1:
                    self.position = np.array(position, 'f')
                    self.n = len(position)
                    break
            elif start_read:
                self.position = np.array(position, 'f')
                self.n = len(position)
                break

    def write_gjf(self, coordinates, chk_path, gjf_path):
        """
        根据坐标构建gjf文件
        """
        lines = []
        lines.append('%chk='+chk_path+'\n')
        lines.append('%nprocshared=4'+'\n')
        lines.append('%mem=1000MB'+'\n')
        lines.append(self.opt+'\n')
        lines.append('\n'+'Title Card Required\n'+'\n')
        # 电荷和自旋多重度
        lines.append('0 1\n')
        for i in range(len(self.element)):
            lines.append(' '+self.element[i]+' '+'%.8f' % coordinates[i*3]+' '+'%.8f' % coordinates[i*3+1]+' '+'%.8f' % coordinates[i*3+2]+'\n')
        lines.append('\n')
        with open(gjf_path, 'w') as f:
            f.writelines(lines)

    def run_gaussian(self, path, filename):
        """
        运行gaussian计算
        """
        os.system('g09 ' + path + '\\' + filename + ' ' + path + '\\' + filename)

    def check_normal_end(self, line):
        """
        检查计算是否正常结束
        """
        result = re.match(line, self.lines[-1])
        if result:
            return True
        else:
            return False

    def write_log(self, info):
        """
        将信息写入日志
        """
        current_time = time.strftime('%Y-%m-%d %H:%M:%S', time.localtime(time.time()))
        text = current_time + ' : ' + info
        print(text)
        with open(self.log_path, 'a') as f:
            f.writelines(text+'\n')

    def calculate(self, coordinates, path, chk_name, gjf_name):
        """
        由坐标向量计算能量和受力
        """
        self.write_gjf(coordinates, path+'\\'+chk_name+'.chk', path+'\\'+gjf_name+'.gjf')
        self.run_gaussian(path, gjf_name)

    def dimer_info(self, coordinates1, coordinates2, number, rotate_times):
        """
        计算双子一次迭代过程中的能量和受力
        每次迭代过程都在一个文件夹下
        """
        # 文件夹不存在则创建文件夹
        if not os.path.exists(self.path+'\\'+str(number)):
            os.mkdir(self.path+'\\'+str(number))
        # 创建输入文件并计算
        name = 'A_' + str(rotate_times)
        self.calculate(coordinates1, self.path+'\\'+str(number), name, name)
        # 计算完成后检查并读取数据
        self.load_file(self.path+'\\'+str(number)+'\\'+name+'.out')
        energy_A = self.search_energy()
        force_A = self.search_force()
        name = 'B_' + str(rotate_times)
        self.calculate(coordinates2, self.path+'\\'+str(number), name, name)
        self.load_file(self.path+'\\'+str(number)+'\\'+name+'.out')
        energy_B = self.search_energy()
        force_B = self.search_force()
        return [energy_A, energy_B, force_A, force_B]


if __name__ == "__main__":
    path = r'C=C=CC_scf.LOG'
    p = PES()
    p.init_load('C=C=CC.gjf')
    p.calculate(p.position, p.path, 'test', 'test')
    # print(p.position)
    # p.log_path = r'temp_log.txt'
    # p.write_log('info test')
    # p.run_gaussian(r'1', 'C=C=CC')
